//
//  scorer.cpp
//  yuzon
//
//  Created by Alvin on 2023/3/29.
//

#include "scorer.hpp"
#include "utils.hpp"

using namespace cv;
using namespace std;

Scorer::Scorer(const string &angle_list_path) :
start_angle_(kStartAngle),
frame_offsets_(kFrameOffsets) {
  angles_ = AngleList(angle_list_path);
}

void Scorer::Process(const KeypointResults &results,
                     int frame,
                     int &state,
                     vector<vector<int>> &frame_scores) {
  switch (state) {
    case -2: { // 等待开始
      if (results.empty())
        break;
      
      const KeypointResult &result = results[0];
      int num_errors = CheckKeypoints(result.keypoints);
      
      if (num_errors > 0)
        break;
      
      vector<float> user_angle = GetUserAngle(result.keypoints);
      int start_score = CalculateScore(user_angle, start_angle_);
      
      // 开始计分
      if (start_score >= 85) {
        state = -1;
        break;
      }
      
      break;
    }
      
    case -1: // 开始计分
    case 1: // 计分异常
      state = 0;
      
    case 0: { // 正常计分
      if (frame < 0)
        break;
      
      if (frame >= angles_.size() - 1) {
        state = 2;
        break;
      }
      
      vector<vector<float>> standard_angles = GetStandardAngles(frame);
      
      // 标准角度为空，不计分
      if (standard_angles.empty())
        break;
      
      vector<int> frame_score{frame};
      
      if (results.empty()) {
        frame_score.emplace_back(1);
        frame_scores.emplace_back(frame_score);
        state = 1;
        break;
      }
      
      const KeypointResult &result = results[0];
      int num_errors = CheckKeypoints(result.keypoints);
      
      if (num_errors > 0) {
        if (num_errors >= 7)
          frame_score.emplace_back(1);
        else {
          frame_score.emplace_back(0);
        }
        
        frame_scores.emplace_back(frame_score);
        state = 1;
        break;
      }
      
      vector<float> user_angle = GetUserAngle(result.keypoints);
      int max_score = CalculateMaxScore(user_angle, standard_angles);
      
      frame_score.emplace_back(max_score);
      frame_scores.emplace_back(frame_score);
      break;
    }
      
    case 2: // 计分结束
      break;
  }
}

int Scorer::CalculateMaxScore(const vector<float> &user_angle,
                              const vector<vector<float>> &standard_angles) {
  int max_score = 0;
  
  for (const vector<float> &standard_angle : standard_angles) {
    int score = CalculateScore(user_angle, standard_angle);
    max_score = max(max_score, score);
  }
  
  return max_score;
}

int Scorer::CalculateScore(const vector<float> &user_angle,
                           const vector<float> &standard_angle) {
  int score = 0;
  
  for (int i = 0; i < 10; ++i) {
    float u = user_angle[i];
    float s = standard_angle[i];
    int n = int(abs(u - s) / 5);
    score += n > 6 ? 1 : 100 - (5 * n);
  }
  
  return round(score / 10);
}

int Scorer::CheckKeypoints(const vector<float> &keypoints) {
  int num_errors = 0;
  
  for (int i = 5; i < keypoints.size() / 3; ++i) {
    if (keypoints[i * 3] >= kKeypointThreshold &&
        keypoints[i * 3 + 1] > 0.0f &&
        keypoints[i * 3 + 2] > 0.0f)
      continue;
    
    ++num_errors;
  }
  
  return num_errors;
}

float Scorer::GetAngle(const Point2f &a, const Point2f &b) {
  Point2f s = b - a;
  Point2f t(0.0f, 1.0f);
  double c = max(-1.0, min(1.0, s.dot(t) / (norm(s) * norm(t))));
  return acos(c) * 180 / M_PI;
}

vector<vector<float>> Scorer::GetStandardAngles(int frame) {
  vector<vector<float>> standard_angles;
  
  for (const int frame_offset : frame_offsets_) {
    int index = frame + frame_offset;
    
    // 超出标准角度范围则跳过
    if (index < 0 || angles_.size() <= index)
      continue;
    
    vector<float> standard_angle = angles_[index];
    
    // 最大角度小于10则跳过
    if (standard_angle[10] < 10.0f)
      continue;
    
    // 取10个标准角度
    standard_angle.pop_back();
    standard_angles.emplace_back(standard_angle);
  }
  
  return standard_angles;
}

vector<float> Scorer::GetUserAngle(const vector<float> &keypoints) {
  // 角度1 左肩和左肘的连线与y方向的夹角 (5, 7)
  float angle_0 = GetAngle(Point2f(keypoints[(5 * 3) + 1],
                                   keypoints[(5 * 3) + 2]),
                           Point2f(keypoints[(7 * 3) + 1],
                                   keypoints[(7 * 3) + 2]));
  
  // 角度2 左肘和左腕的连线与y方向的夹角 (7, 9)
  float angle_1 = GetAngle(Point2f(keypoints[(7 * 3) + 1],
                                   keypoints[(7 * 3) + 2]),
                           Point2f(keypoints[(9 * 3) + 1],
                                   keypoints[(9 * 3) + 2]));
  
  // 角度3 右肩和右肘的连线与y方向的夹角 (6, 8)
  float angle_2 = GetAngle(Point2f(keypoints[(6 * 3) + 1],
                                   keypoints[(6 * 3) + 2]),
                           Point2f(keypoints[(8 * 3) + 1],
                                   keypoints[(8 * 3) + 2]));
  
  // 角度4 右肘和右腕的连线与y方向的夹角 (8, 10)
  float angle_3 = GetAngle(Point2f(keypoints[(8 * 3) + 1],
                                   keypoints[(8 * 3) + 2]),
                           Point2f(keypoints[(10 * 3) + 1],
                                   keypoints[(10 * 3) + 2]));
  
  // 角度5 左胯和左膝的连线与y方向的夹角 (11, 13)
  float angle_4 = GetAngle(Point2f(keypoints[(11 * 3) + 1],
                                   keypoints[(11 * 3) + 2]),
                           Point2f(keypoints[(13 * 3) + 1],
                                   keypoints[(13 * 3) + 2]));
  
  // 角度6 左膝和左踝的连线与y方向的夹角 (13, 15)
  float angle_5 = GetAngle(Point2f(keypoints[(13 * 3) + 1],
                                   keypoints[(13 * 3) + 2]),
                           Point2f(keypoints[(15 * 3) + 1],
                                   keypoints[(15 * 3) + 2]));
  
  // 角度7 右胯和右膝的连线与y方向的夹角 (12, 14)
  float angle_6 = GetAngle(Point2f(keypoints[(12 * 3) + 1],
                                   keypoints[(12 * 3) + 2]),
                           Point2f(keypoints[(14 * 3) + 1],
                                   keypoints[(14 * 3) + 2]));
  
  // 角度8 右膝和右踝的连线与y方向的夹角 (14, 16)
  float angle_7 = GetAngle(Point2f(keypoints[(14 * 3) + 1],
                                   keypoints[(14 * 3) + 2]),
                           Point2f(keypoints[(16 * 3) + 1],
                                   keypoints[(16 * 3) + 2]));
  
  // 角度9 左肩和左胯的连线与y方向的夹角 (5, 11)
  float angle_8 = GetAngle(Point2f(keypoints[(5 * 3) + 1],
                                   keypoints[(5 * 3) + 2]),
                           Point2f(keypoints[(11 * 3) + 1],
                                   keypoints[(11 * 3) + 2]));
  
  // 角度10 右肩和右胯的连线与y方向的夹角 (6, 12)
  float angle_9 = GetAngle(Point2f(keypoints[(6 * 3) + 1],
                                   keypoints[(6 * 3) + 2]),
                           Point2f(keypoints[(12 * 3) + 1],
                                   keypoints[(12 * 3) + 2]));
  
  vector<float> user_angle{
    angle_0, angle_1, angle_2, angle_3, angle_4,
    angle_5, angle_6, angle_7, angle_8, angle_9
  };
  
  return user_angle;
}
